.text
.global task_switch
.extern current_task

// void task_switch(void* next);
// 
// 参数:
//   x0 - 指向下一个任务的TCB指针
//
// 注意: 在AArch64调用约定中:
// - 调用者保存寄存器: x0-x18
// - 被调用者保存寄存器: x19-x29, LR (x30)
//
// 我们只需要保存被调用者保存寄存器

task_switch:
    // 获取当前任务的TCB指针
    adrp x9, current_task
    add x9, x9, :lo12:current_task
    ldr x10, [x9]
    
    // 检查current_task是否为NULL
    cbz x10, skip_save  // 如果为空，跳过保存当前上下文
    
    // 保存当前上下文到当前任务的栈上
    stp x29, x30, [sp, #-16]!  // 保存帧指针和返回地址
    stp x19, x20, [sp, #-16]!
    stp x21, x22, [sp, #-16]!
    stp x23, x24, [sp, #-16]!
    stp x25, x26, [sp, #-16]!
    stp x27, x28, [sp, #-16]!
    
    // 保存当前栈指针到当前任务的TCB中
    mov x11, sp         // 使用临时寄存器x11存储sp的值
    str x11, [x10]      // 将栈指针值存储到TCB
    
skip_save:
    // 切换到新任务
    str x0, [x9]  // x0是第一个参数，即下一个任务的TCB指针
    
    // 加载新任务的栈指针
    ldr x11, [x0]       // 从TCB加载栈指针到临时寄存器
    mov sp, x11         // 设置栈指针
    
    // 恢复新任务的上下文
    ldp x27, x28, [sp], #16
    ldp x25, x26, [sp], #16
    ldp x23, x24, [sp], #16
    ldp x21, x22, [sp], #16
    ldp x19, x20, [sp], #16
    ldp x29, x30, [sp], #16  // 恢复帧指针和返回地址
    
    // 返回到新任务
    ret
